import logging
import os

import torch

from PIL import Image
from torchvision.transforms.functional import to_tensor, to_pil_image

from models.DnCNN_noiseprint import make_net

num_levels = 17
out_channel = 1
noiseprint = make_net(3, kernels=[3, ] * num_levels,
                           features=[64, ] * (num_levels - 1) + [out_channel],
                           bns=[False, ] + [True, ] * (num_levels - 2) + [False, ],
                           acts=['relu', ] * (num_levels - 1) + ['linear', ],
                           dilats=[1, ] * num_levels,
                           bn_momentum=0.1, padding=1)

noiseprint_path = "/home/wc/disk1/MMFusion/pretrained/noiseprint/np++.pth"
if noiseprint_path:
    np_weights = noiseprint_path
    assert os.path.isfile(np_weights)
    dat = torch.load(np_weights, map_location=torch.device('cpu'))
    logging.info(f'Noiseprint++ weights: {np_weights}')
    noiseprint.load_state_dict(dat)

noiseprint.eval()
for param in noiseprint.parameters():
    param.requires_grad = False

if __name__ == '__main__':
    img = Image.open('/home/wc/disk1/datasets/COVER/tampered/1t.tif').convert('RGB')
    import torchvision.transforms.functional as TF

    inp = to_tensor(img).unsqueeze(0)
    images_norm = TF.normalize(to_tensor(img), mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    import matplotlib.pyplot as plt
    out = noiseprint(inp)
    if out.size()[-3] == 1:
        out = torch.tile(out, (3, 1, 1))
    # out_norm = TF.normalize(out, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

    import numpy as np
    fig, ax = plt.subplots(1, 2)
    # noiseprint = out[:, 0].squeeze().numpy()
    # # out = torch.sigmoid(out)
    ax[0].imshow(img)
    ax[0].set_title('Image')
    # #
    #
    noiseprint = out[:,0].squeeze().numpy()
    ax[1].imshow(noiseprint[16:-16:4, 16:-16:4], cmap='gray')
    # ax[1].imshow(out, cmap='gray')
    ax[1].set_title('NoisePrint++')

    plt.show()